import sys
if not sys.platform.startswith('win'):
    from deepmd.infer import DeepPot
from core_atomistic.atomic_model import AtomicModel
import numpy as np


class DeepmdCalc(object):
    """The atom class."""

    def __init__(self):
        """Constructor"""
        self.pot: DeepPot = None
        self.type_map = {}

    def set_potential(self, filename, t_map=dict([(3, 0), (6, 1)])):
        self.pot = DeepPot(filename)
        self.type_map = t_map

    def potential(self, model: AtomicModel):
        if self.pot is None:
            print("No potential")
            return

        coord = model.get_positions().reshape([1, -1])
        cell = model.lat_vectors.reshape([1, -1])
        ch = model.get_atomic_numbers()
        ch = np.array(ch, dtype=int)
        atype = np.zeros(len(ch), dtype=int)
        for i in range(len(ch)):
            atype[i] = self.type_map[ch[i]]

        mixed_type = False
        e, f, v = self.pot.eval(coord, cell, atype, mixed_type=mixed_type)
        # where e, f and v are predicted energy, force and virial of the system, respectively.
        return e[0]
